package nl.cypherpunk.proLearner.rtsp;


import java.io.*;
import java.net.*;
import java.nio.ByteBuffer;
import java.text.DecimalFormat;
import java.util.*;
import java.util.concurrent.TimeoutException;

import static nl.cypherpunk.proLearner.Common.intToByteArray;

public class RTSPTestService {
    Socket socket;
    OutputStream output;
    InputStream input;

    String host = "127.0.0.1";
    int port = 8554;
    String hostname = "localhost";
    // Act as a TLS client
    boolean ROLE_CLIENT = true;
    // Restart server after every session
    boolean REQUIRE_RESTART = true;
    // Timeout in ms
    int RECEIVE_MSG_TIMEOUT = 50;
    // Send output from TLS implementation to console
    boolean CONSOLE_OUTPUT = false;

    String cmd;
    public Process targetProcess;

    String redirectFile;
    ArrayList<byte[]> currentCase;
    int totalExecs;
    int prevCoverage;
    String outputDir;
    String env_cmd="";
    PrintWriter out;
    BitSet coverageMap;
    Double initScore;
    int seq;
    String session;
    private static final String RTSP_OK = "RTSP/1.0";
    private static final String HTTP_OK = "HTTP/1.0";
    String trackInfo = "1";
    HashSet<String> crashHash;
    ArrayList<String> symbols;
    Map<String,Integer> subtypes;
    LinkedHashMap<String, byte[]> mapper;
    HashSet<ByteBuffer> messages;
    public RTSPTestService() throws Exception {
        symbols = new ArrayList<>();
        mapper = new LinkedHashMap<>();
        subtypes = new LinkedHashMap<>();
        messages = new HashSet<>();
    }

    public void setTarget(String target) throws Exception {
        if(target.equals("server")) {
            ROLE_CLIENT = true;
        }
        else if(target.equals("client")) {
            ROLE_CLIENT = false;
        }
        else {
            throw new Exception("Unknown target");
        }
    }

    public void setHost(String host) {
        this.host = host;
    }

    public void setPort(int port) {
        this.port = port;
    }

    public void setCommand(String cmd) {
        this.cmd = cmd;
    }

    public void setRestartTarget(boolean restart) {
        this.REQUIRE_RESTART = restart;
    }

    public void setReceiveMessagesTimeout(int timeout) {
        RECEIVE_MSG_TIMEOUT = timeout;
    }

    public void setConsoleOutput(boolean enable) {
        CONSOLE_OUTPUT = enable;
    }

    public void setOutputDir(String dir){
        this.outputDir = dir;
    }
    public void setRedirectFile(String output) {
        redirectFile = output;
    }

    public void saveMapper()
    {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        for(Map.Entry<String, byte[]> entry : mapper.entrySet()) {
            try {
                file = new File(this.outputDir + File.separator + entry.getKey());
                file.createNewFile();
                fos = new FileOutputStream(file);
                bos = new BufferedOutputStream(fos);
                byte[] tmp = intToByteArray(entry.getValue().length);
                bos.write(tmp, 0, 4);
                bos.write(entry.getValue(), 0, entry.getValue().length);

            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                if (bos != null) {
                    try {
                        bos.close();
                    } catch (IOException e1) {
                        e1.printStackTrace();
                    }
                }
                if (fos != null) {
                    try {
                        fos.close();
                    } catch (IOException e1) {
                        e1.printStackTrace();
                    }
                }
            }
        }
    }

    public void start() throws Exception {
        currentCase = new ArrayList<>();
        prevCoverage = 0;
        totalExecs = 0;
        crashHash = new HashSet<>();
        if(ROLE_CLIENT) {

            if(cmd != null && !cmd.equals("")) {
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //pb.redirectOutput(new File(redirectFile));
                targetProcess = pb.start();
                Thread.sleep(200);
            }
            //connectSocket();
            session = "";
            seq = 1;
        }
        else {
            RTSPTestService.RTSPTestServiceRunnable rtspTestService = this.new RTSPTestServiceRunnable(this);
            rtspTestService.start();
            Thread.sleep(50);
            if(cmd != null && !cmd.equals("")) {
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //pb.redirectOutput(new File(redirectFile));
                Map< String,String> env = pb.environment();
                env.put("LD_PRELOAD",env_cmd);

				if(CONSOLE_OUTPUT) {
					pb.inheritIO();
				}
				/*
				else {
					pb.redirectErrorStream(true);
					pb.redirectOutput(new File("output.log"));
				}
				*/
                targetProcess = pb.start();
                //tlsClient = new TLSClient(targetProcess);
            }

            // Wait for the client to send the first message (ClientHello)
            while(!rtspTestService.isReady()) Thread.sleep(10);

        }
    }
    public void reset() throws Exception {
        //System.out.println("RESET");
        if(socket!=null)
            socket.close();
        totalExecs ++;
        if(totalExecs%1000==0)
            System.out.println(totalExecs);
        if(ROLE_CLIENT) {
            if (!currentCase.isEmpty()) {
                int coverage = getCoverage();
                //reader.close();
                if (coverage == 1)
                    seedSave(currentCase, this.outputDir + File.separator + "Crash" + File.separator + "Crash_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));
                //if (coverage == 2)
                //    seedSave(currentCase, this.outputDir + File.separator + "Seed" + File.separator + "Abnormal_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));
                if (coverage > prevCoverage) {
                    System.out.println(coverage);
                    seedSave(currentCase, this.outputDir + File.separator + "Seed" + File.separator + "A_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));

                    prevCoverage = coverage;
                }
                currentCase.clear();
            }
            if (REQUIRE_RESTART && cmd != null && !cmd.equals("")) {
                targetProcess.destroy();
                //String target = this.outputDir.substring(0, this.outputDir.indexOf("2"));
                String killall = "sudo killall testOnDemandRTSPServer";
                if (REQUIRE_RESTART) {
                    ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
                    Process tmp = pbkill.start();
                    tmp.waitFor();
                }
                Thread.sleep(100);
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //File log = new File("log.txt");
                //pb.redirectOutput(new File(redirectFile));
                targetProcess = pb.start();

                Thread.sleep(100);
            }
            //Thread.sleep(100);
            try {
                connectSocket();
            }
            catch (Exception e)
            {
                String killall = "sudo killall testOnDemandRTSPServer";
                if (REQUIRE_RESTART) {
                    ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
                    Process tmp = pbkill.start();
                    tmp.waitFor();
                }
                Thread.sleep(200);
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //File log = new File("log.txt");
                //pb.redirectOutput(new File(redirectFile));
                targetProcess = pb.start();

                Thread.sleep(500);
                connectSocket();
            }
            seq = 1;
            session = "";
        } else {
            if (!currentCase.isEmpty()) {
                int coverage = getCoverage();
                //reader.close();
                if (coverage == 1)
                    seedSave(currentCase, this.outputDir + File.separator + "Crash" + File.separator + "Crash_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));
                //if (coverage == 2)
                //    seedSave(currentCase, this.outputDir + File.separator + "Seed" + File.separator + "Abnormal_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));
                if (coverage > prevCoverage) {
                    System.out.println(coverage);
                    seedSave(currentCase, this.outputDir + File.separator + "Seed" + File.separator + "A_" + String.format("%06d", totalExecs) + "_" + String.format("%05d", coverage));

                    prevCoverage = coverage;
                }
                currentCase.clear();
            }
            String killall = "sudo killall testRTSPClient";
            ProcessBuilder pbkill = new ProcessBuilder(killall.split(" "));
            Process tmp = pbkill.start();
            tmp.waitFor();
            Thread.sleep(50);
            RTSPTestService.RTSPTestServiceRunnable rtspTestService = this.new RTSPTestServiceRunnable(this);
            rtspTestService.start();
            Thread.sleep(50);
            if(cmd != null && !cmd.equals("")) {
                ProcessBuilder pb = new ProcessBuilder(cmd.split(" "));
                pb.redirectErrorStream(true);
                //File log = new File(redirectFile);
                //pb.redirectOutput(new File(redirectFile));
                Map< String,String> env = pb.environment();
                env.put("LD_PRELOAD",env_cmd);

				if(CONSOLE_OUTPUT) {
					pb.inheritIO();
				}
				/*
				else {
					pb.redirectErrorStream(true);
					pb.redirectOutcrashHashput(new File("output.log"));
				}
				*/
                targetProcess = pb.start();
                Thread.sleep(50);

                session = "";
                receiveMessages();
                //tlsClient = new TLSClient(targetProcess);
            }

            // Wait for the client to send the first message (ClientHello)

        }
    }

    public String receiveMessages() throws Exception {
        String out = "";
        String statusCode = "";
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        if(socket.isInputShutdown())
            return "ConnectionClosed";
        try {
            byte[] buffer = new byte[1024];
            int len = 0;
            while((len = input.read(buffer)) != -1) {
                if(len==0)
                    break;
                bos.write(buffer, 0, len);
                //if(bos.toString().endsWith("\r\n\r\n"))
                //    break;
            }
            bos.close();
            out = bos.toString();
        }
        catch (SocketTimeoutException e)
        {
            //socket.close();
            bos.close();
            if(bos.size()==0)
                out = "Empty";
            else
                out = bos.toString();
        }
        String[] tmp = out.split("\r\n\r\nRTSP");
        for (int i = 0; i < tmp.length; i++) {
            String reponse = "";
            if(i==0)
                reponse = tmp[0];
            else
                reponse = "RTSP"+tmp[i];
            if (reponse.startsWith(RTSP_OK)) {
                statusCode = statusCode + reponse.substring(reponse.indexOf(RTSP_OK)+9,reponse.indexOf(RTSP_OK)+12);
            } else if (reponse.startsWith(HTTP_OK)) {
                statusCode = statusCode + reponse.substring(reponse.indexOf(HTTP_OK)+9,reponse.indexOf(HTTP_OK)+12);
            }
            else if (reponse.startsWith("SETUP")) {
                if(reponse.contains("track1"))
                    statusCode = statusCode + "SETUP1";
                else if(reponse.contains("track2"))
                    statusCode = statusCode + "SETUP2";
            } else if (reponse.startsWith("PLAY")) {
                statusCode = statusCode + "PLAY";
            } else if (reponse.startsWith("TEARDOWN")) {
                statusCode = statusCode + "TEARDOWN";
            } else if (reponse.startsWith("DESCRIBE")) {
                statusCode = statusCode + "DESCRIBE";
            }
            else {
                //if(reponse.length() > 4)
                //    System.out.println(reponse.substring(0,4));
                statusCode = "ConnectionClosed";
            }

            if(statusCode.contains("ConnectionClosed") && ROLE_CLIENT) {
                try {
                    if (!socket.isOutputShutdown()) {
                        socket.close();
                    }
                } catch (Exception e)
                {
                    socket.close();
                }
            }
        }

        return statusCode;
    }

    //update mapper and subtypes
    public boolean addMapper(String type, byte[] message){
        String symbol="";
        if(messages.contains(ByteBuffer.wrap(message)))
            return false;
        messages.add(ByteBuffer.wrap(message));
        if(subtypes.containsKey(type))
        {
            Integer index = subtypes.get(type)+1;
            subtypes.put(type,index);
            symbol=type+(new DecimalFormat("00").format(index));
        }
        else
        {
            subtypes.put(type,1);
            symbol = type+"01";
        }
        symbols.add(symbol);
        mapper.put(symbol,message);
        return true;
    }
    //reverse order
    public boolean delMapper(String subtype){
        symbols.remove(subtype);
        byte[] tmp = mapper.remove(subtype);
        messages.remove(ByteBuffer.wrap(tmp));
        String type = subtype.substring(0,subtype.length()-2);
        if(subtypes.get(type)==1)
        {
            subtypes.remove(type);
        }
        else
        {
            subtypes.put(type,subtypes.get(type)-1);
        }
        return true;
    }

    public String processSymbol(String input) throws Exception {
        String inAction = input;
        byte[] out = null;
        if(!socket.isConnected() || socket.isClosed()) return "ConnectionClosed";

        try {
            out = mapper.get(inAction);
            sendMessage(out);
            currentCase.add(out);
            seq ++;
            return receiveMessages();
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }
    }

    public String buildTeardown() throws Exception {
        String message = "TEARDOWN rtsp://"+host+":"+Integer.toString(port)+"/matroskaFileTest/ RTSP/1.0\r\n";
        message = message.concat("CSeq: "+Integer.toString(seq)+"\r\n");
        message = message.concat("User-Agent: ./testRTSPClient (LIVE555 Streaming Media v2018.08.28)\r\n");
        message = message.concat("Session: "+session+"\r\n");
        return message;
    }

    public String sendFuzzingMessage (byte[] message) throws Exception
    {
        try {
            if (socket.isClosed())
                return "ConnectionClosed";
            output.write(message);
            return receiveFuzzingMessages();
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }
    }

    public String receiveFuzzingMessages() throws Exception {
        try{
            String out = receiveMessages();
            if(out.contains("ConnectionClosed")) {
                socket.close();
            }
            return out;
        }
        catch(SocketException e) {
            //String outAction = "ConnectionClosedException";
            String outAction = "ConnectionClosed";

            return outAction;
        }

    }

    public String connectSocket() throws UnknownHostException, IOException,Exception {
        socket = new Socket(host, port);
        socket.setTcpNoDelay(true);
        socket.setSoTimeout(RECEIVE_MSG_TIMEOUT);
        socket.setKeepAlive(false);
        socket.setReuseAddress(true);
        output = socket.getOutputStream();
        input = socket.getInputStream();
        return "";
    }

    public void listenSocket() throws UnknownHostException, IOException {
        ServerSocket server = new ServerSocket();
        server.bind(new InetSocketAddress(host, port));
        socket = server.accept();
        socket.setTcpNoDelay(true);
        socket.setSoTimeout(RECEIVE_MSG_TIMEOUT);

        output = socket.getOutputStream();
        input = socket.getInputStream();

        server.close();
    }

    public void closeSocket() throws IOException, Exception {
        socket.close();
    }
    void sendMessage(byte[] msg) throws Exception {
        if(!socket.isOutputShutdown())
            output.write(msg);
    }

    public void close() {
        if(targetProcess != null) {
            targetProcess.destroy();
        }
    }

    class RTSPTestServiceRunnable extends Thread {
        RTSPTestService rtsp;
        boolean ready;

        public RTSPTestServiceRunnable(RTSPTestService rtsp) {
            ready = false;
            this.rtsp = rtsp;
        }

        public boolean isReady() {
            return ready;
        }

        public boolean isConnected() {
            return rtsp.socket.isConnected();
        }

        public boolean isBound() {
            return  (rtsp.socket != null) && rtsp.socket.isBound();
        }

        public void run() {
            try {
                rtsp.listenSocket();
                rtsp.receiveMessages();
                ready = true;
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    //change by pany
    public void seedSave(List<byte[]> seed, String filename) {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        try {
            file = new File(filename);
            if (file.exists()) {
                return;
            }
            fos = new FileOutputStream(file);
            bos = new BufferedOutputStream(fos);
            for (int i = 0; i < seed.size(); i++) {
                byte[] tmp = intToByteArray(seed.get(i).length);
                bos.write(tmp, 0, 4);
                bos.write(seed.get(i), 0, seed.get(i).length);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (bos != null) {
                try {
                    bos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }
    }

    public int getCoverage() throws IOException {
        int coverage = 0;
        String resultStr = new String();
        /*if (ROLE_CLIENT) {
            //BufferedReader reader = new BufferedReader(new InputStreamReader(targetProcess.getInputStream()));
            StringBuilder builder = new StringBuilder();
            String line = null;

            while (reader.ready()) {
                line = reader.readLine();
                if (line == null)
                    break;
                if (line.contains("Coverage")) {
                    try {
                        int index = line.indexOf("Coverage");
                        String cover = line.substring(index + 9).trim();
                        coverage = Integer.parseInt(cover);
                        //break;
                    } catch (NumberFormatException e) {
                        e.printStackTrace();
                    }
                }
                if (line.contains("AddressSanitizer"))
                    coverage = 1;
                resultStr = resultStr.concat(line);
            }
        } else {*/
        if (coverageMap == null)
            coverage = 0;
        else
            coverageMap.cardinality();
        BufferedReader reader = new BufferedReader(new InputStreamReader(targetProcess.getInputStream()));
        StringBuilder builder = new StringBuilder();
        String line = null;
        try {
            //Thread.sleep(100);
            while (reader.ready()) {
                line = reader.readLine();
                if (line == null)
                    break;
                if (line.contains("Bitmap")) {
                    try {
                        int index = line.indexOf("Bitmap");
                        String cover = line.substring(index + 8).trim();
                        if (coverageMap == null) {
                            coverageMap = new BitSet(cover.length());
                        }
                        for (int i = 0; i < cover.length(); i++)
                            if (cover.charAt(i) == '1')
                                coverageMap.set(i);
                        coverage = coverageMap.cardinality();
                        //System.out.print(line);
                        //break;
                    } catch (NumberFormatException e) {
                        e.printStackTrace();
                    }
                }
                if (line.contains("AddressSanitizer") && !line.contains("SUMMARY: AddressSanitizer:")) {
                    //System.out.println(line);
                    Thread.sleep(500);
                }
                if (line.contains("SUMMARY: AddressSanitizer:")) {
                    coverage = 1;
                    //System.out.println(line);
                    if (!crashHash.add(line))
                        return 2;
                    resultStr = resultStr.concat(line + "\r\n");
                    System.out.println(resultStr);
                    return coverage;
                }
                if (line.contains("Aborted")) {
                    coverage = 1;
                    if (!crashHash.add(line))
                        return 2;
                    System.out.println(line);
                    return coverage;
                }
                resultStr = resultStr.concat(line + "\r\n");
                //System.out.print(line);
            }
        } catch (Exception e) {
            System.out.println(e);
        }
        //}
        if (resultStr.length() == 0)
            coverage = 3;
        //System.out.print(resultStr);
        if (ROLE_CLIENT){
            if (targetProcess.isAlive())
                return coverage;
            else {
                System.out.println("Aborted");
                return 1;
            }
        }
        return coverage;
    }
}
